import os
import argparse

import torch
import matplotlib.pyplot as plt
import matplotlib.patches as patches

from test import GOTURN

import cv2
import numpy as np
from shutil import copyfile

import glob as gb
import shutil
import picture_frame

args = None
parser = argparse.ArgumentParser(description='GOTURN Testing')
parser.add_argument('-w', '--model-weights',default='../models/final_model.pth',
                    type=str, help='path to pretrained model')
parser.add_argument('-d', '--data-directory',
                    default='../data/OTB/test', type=str,
                    help='path to video frames')
parser.add_argument('-s', '--save-directory',
                    default='../result',
                    type=str, help='path to save directory')
parser.add_argument('-i', '--input-directory',
                    type=str, help='path to avi directory')

def axis_aligned_iou(boxA, boxB):
    # make sure that x1,y1,x2,y2 of a box are valid
    assert(boxA[0] <= boxA[2])
    assert(boxA[1] <= boxA[3])
    assert(boxB[0] <= boxB[2])
    assert(boxB[1] <= boxB[3])

    # determine the (x, y)-coordinates of the intersection rectangle
    xA = max(boxA[0], boxB[0])
    yA = max(boxA[1], boxB[1])
    xB = min(boxA[2], boxB[2])
    yB = min(boxA[3], boxB[3])

    # compute the area of intersection rectangle
    interArea = max(0, xB - xA + 1) * max(0, yB - yA + 1)

    # compute the area of both the prediction and ground-truth
    # rectangles
    boxAArea = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1)
    boxBArea = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1)

    # compute the intersection over union by taking the intersection
    # area and dividing it by the sum of prediction + ground-truth
    # areas - the interesection area
    iou = interArea / float(boxAArea + boxBArea - interArea)

    # return the intersection over union value
    return iou


def save(ax, im, bb, gt_bb, idx):
    ax.clear()
    ax.imshow(im)
    goturn_box = patches.Rectangle((bb[0], bb[1]),
                                   bb[2]-bb[0], bb[3]-bb[1],
                                   linewidth=2, edgecolor='r',
                                   facecolor='none')
    gt_box = patches.Rectangle((gt_bb[0], gt_bb[1]),
                               gt_bb[2]-gt_bb[0], gt_bb[3]-gt_bb[1],
                               linewidth=2, edgecolor='w',
                               facecolor='none')
    ax.add_patch(goturn_box)
    ax.add_patch(gt_box)
    props = dict(boxstyle='round', facecolor='red', alpha=0.5)
   # ax.text(0.05, 0.95, "GOTURN", transform=ax.transAxes, fontsize=12,
   #         verticalalignment='top', bbox=props)
    ax.set_axis_off()
    plt.savefig(os.path.join(args.save_directory, str(idx)+'.jpg'))


def main(args,pic_num):
    use_gpu = torch.cuda.is_available()
    tester = GOTURN('../data/OTB/'+str(pic_num-1),
                    args.model_weights,
                    use_gpu)
    fig, ax = plt.subplots(1)
    if os.path.exists(args.save_directory):
        print('Save directory %s already exists' % (args.save_directory))
    else:
        os.makedirs(args.save_directory)
    # save initial frame with bounding box
    #save(ax, tester.img[0][0], tester.prev_rect, tester.prev_rect, 1)
    tester.model.eval()

    # loop through sequence images
    for i in range(tester.len):
        # get torch input tensor
        sample = tester[i]

        # predict box
        bb = tester.get_rect(sample)
        gt_bb = tester.gt[i]
        tester.prev_rect = bb

        # save current image with predicted rectangle and gt box
        im = tester.img[i][1]
        save(ax, im, bb, gt_bb, i+1+(pic_num-1)*500)

        # print stats
        print('frame: %d, IoU = %f' % (
            i+2, axis_aligned_iou(gt_bb, bb)))

def init_picture(args):
    del_file('../result')
    del_file('../data/OTB')
    pic_res=0 #是否选定检测帧
    pic_num=0
    pathName=args.input_directory
    while pic_res==0:
        pic_count=0 #视频帧数统计
        capture = cv2.VideoCapture(pathName)
        while(capture.isOpened()):
            ret, frame = capture.read()
            if ret == True:
                frame=cv2.resize(frame,(640,480),interpolation=cv2.INTER_CUBIC)
                if pic_res==0:
                    cv2.imshow('frame', frame)
                    pic_count=pic_count+1
                    if cv2.waitKey(40) & 0xFF == ord('q'):    #点击视频窗口，按q键选定图片
                        pic_res =1
                        frame_res=frame 
                        img=frame_res
                        (a,b)=picture_frame.get_rect(img,title='get_rect') 
                        coordinate=str(a[0])+','+str(a[1])+','+str(b[0]-a[0])+','+str(b[1]-a[1])              
                        cv2.destroyAllWindows()
                else:
                    pic_count=pic_count+1
            else:
                break
        capture.release()    
    if pic_res!=0:
        if pic_count%500==0:
            pic_num=pic_count/500
        else:
            pic_num=pic_count//500+1
        capture = cv2.VideoCapture(pathName)
        for i in range(0,pic_num):
            img=frame_res
            curPath=os.getcwd()
            tempPath='../data/OTB/'+str(i)+'/img'
            targetPath=curPath+os.path.sep+tempPath
            if not os.path.exists(targetPath):
                os.makedirs(targetPath)
                cv2.imwrite(targetPath+'/0001.jpg', frame)
            else:
                print('路径已经存在！')
            f=open('../data/OTB/'+str(i)+'/groundtruth_rect.txt','w')
            f.write(coordinate)  
            cv2.imwrite('../data/OTB/'+str(i)+'/img/0001.jpg', frame_res)    
            for m in range(1,501):
                ret, frame = capture.read()
                if ret == True:
                    frame=cv2.resize(frame,(640,480),interpolation=cv2.INTER_CUBIC)
                    picture_name='../data/OTB/'+str(i)+'/img/'+str(m+1).zfill(4)+'.jpg'    
                    cv2.imwrite(picture_name, frame)
                    f.write('\n0,0,0,0')
                else:
                    break
        f.close
        capture.release()
    return pic_num

def picture_avi():
    picturecount =0  #计数图片
    picture_number =1 
    videoWriter = cv2.VideoWriter('../result.mp4', cv2.VideoWriter_fourcc(*'MJPG'), 25, (640,480))
    for picturename in os.listdir('../result'):
        picturecount += 1
    for m in range(1,picturecount+1):
        while (os.path.exists("../result/"+str(picture_number)+".jpg")!=True):
            picture_number += 1
        img  = cv2.imread("../result/"+str(picture_number)+".jpg") 
        img = cv2.resize(img,(640,480))
        videoWriter.write(img)
        picture_number += 1

def del_file(path):
    for i in os.listdir(path):
        path_file = os.path.join(path,i)  
        if os.path.isfile(path_file):
            os.remove(path_file)
        else:
            del_file(path_file)

if __name__ == "__main__":
    args = parser.parse_args() 
    pic_num=init_picture(args)
    if pic_num!=0:
        for i in range(1,pic_num+1):
            pic_num=main(args,i)
        picture_avi()
